'''
Created on Jan, 14 2011

@author: oabalbin
'''
import logging
import os
import sys
import subprocess
import tempfile

from optparse import OptionParser
from exome.jobs.base import JOB_SUCCESS, JOB_ERROR
from exome.jobs.job_runner import qsub_cac, qsub_loc, run_local
from exome.jobs.config import ExomePipelineConfig, ExomeAnalysisConfig

# Global Variables
NODE_MEM=45000.0
NODE_CORES=12
SINGLE_CORE=1
MEM_PER_CORE= int(float(NODE_MEM) / NODE_CORES)
# wt=walltime
WT_SHORT= "24:00:00"
WT_LONG= "60:00:00" #"100:00:00"
WT_ELONG="100:00:00"

def gatk_combineVariants(ref_genome, tumor_vcf, benign_vcf, outfile, use_mem, num_cores, 
                         path_to_gatk, merge_option):
    '''
    It combines the variants store in two different VCF files
    merged_options UNION, INTERSECT
    '''
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    
    args = ['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,'-T','CombineVariants',
            '-R', ref_genome, '-variantMergeOptions',merge_option, '-B:tumor,VCF', tumor_vcf,
            '-B:benign,VCF', benign_vcf,'-priority tumor,benign', '-o', outfile]
    
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command

def gatk_VariantEval(ref_genome, merged_tumor_benign_vcf, report_file_sufix, 
                           interesting_snps_vcf, snps_db, use_mem, num_cores,
                           path_to_gatk):
    '''
    It makes an evaluation of the variants according
    to different criteria defined by the jexl_expressions.
    '''
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    
    args = ['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,'-T','VariantEval',
            '-R', ref_genome, '-B:eval,VCF', merged_tumor_benign_vcf, '-B:dbsnp,VCF', snps_db,
            '-reportType', 'R', '-reportLocation', report_file_sufix, '-outputVCF',interesting_snps_vcf,
            '-l','INFO']
    jxel_exp=['-select','set=="Intersection"','-selectName','Intersection', '-select', 'set=="tumor"','-selectName', 'tumor', 
              '-select', 'set=="tumor-filterInbenign"', '-selectName', 'InTumor-FilteredInBenign',
              '-select', 'set=="Benign"', '-selectName', 'benign', 
              '-select', 'set=="benign-filterIntumor"', '-selectName', 'InBenign-FilteredInTumor',
              '-select', 'set=="FilteredInAll"', '-selectName', 'FilteredInAll']
    args=args+jxel_exp
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def gatk_VariantEval2(ref_genome, tumor_vcf, benign_vcf, report_file_sufix, 
                           interesting_snps_vcf, snps_db, use_mem, num_cores,
                           path_to_gatk):
    '''
    It makes an evaluation of the variants according
    to different criteria defined by the jexl_expressions.    
    '''
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    
    
    args = ['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,'-T','VariantEval',
            '-R', ref_genome, '-B:eval,VCF', tumor_vcf, '-B:comp,VCF', benign_vcf, '-B:dbsnp,VCF', snps_db,
            '-reportType', 'R', '-reportLocation', report_file_sufix, '-outputVCF',interesting_snps_vcf,
            '-l','INFO']
    
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command



def gatk_VariantEvaluation(analysis, configrun, num_processors, jobrunfunc, snps_to_use, merge_sample_types=False):
    '''
    It takes calls made on tumor and benign samples and compare them 
    between them and with dbsnp.
    other resources such as hap_map can be included in the comparison
    '''
    
    # Parameters
    extra_mem, num_cores = configrun.gatk_use_mem, configrun.gatk_num_cores
    path_to_gatk, path_to_picard, path_to_sam = configrun.gatk_path, configrun.picard_path, \
                                                configrun.samtools_path
    path_to_Rscript = configrun.rscript_bin
    resources_folder = configrun.gatk_resource_path
    temp_dir = configrun.gatk_temp_dir 
    fdr_filter_level = configrun.gatk_snps_fdr_threshold 
    target_exons = configrun.gatk_target_exons
    genomes = configrun.genomes['human']
    ref_genome, snpdb_vcf, indeldb_file = genomes.gatk_ref_genome, genomes.snpdb, \
                                         genomes.indeldb
    hapmap_vcf, tgk_vcf = genomes.hapmap, genomes.OneKgenomes 
    my_email=configrun.email_addresses
    
    if snps_to_use:
        # use snps found by hard thresholding    
        tumor_vcf= analysis.tumor_snps_hard_vcf
        benign_vfc= analysis.benign_snps_hard_vcf 
        merged_tumor_benign_vcf = analysis.hard_merged_tumor_benign_vcf
        merge_option='UNION'
        variant_eval_file_sufix = analysis.hard_variant_eval_file_sufix
        interesting_snps_vcf = analysis.hard_interesting_snps_vcf
    else:
        tumor_vcf= analysis.tumor_snps_model_vcf
        benign_vfc= analysis.benign_snps_model_vcf
        merged_tumor_benign_vcf = analysis.model_merged_tumor_benign_vcf
        merge_option='UNION'
        variant_eval_file_sufix = analysis.model_variant_eval_file_sufix
        interesting_snps_vcf = analysis.model_interesting_snps_vcf

        
    jobn='veval'
    
    if merge_sample_types:
        command = gatk_combineVariants(ref_genome, tumor_vcf, benign_vfc, merged_tumor_benign_vcf, 
                                       extra_mem, SINGLE_CORE, path_to_gatk, merge_option)
            
        jobidrc = jobrunfunc(jobn+'_mvcf', command, SINGLE_CORE, cwd=None, walltime=WT_SHORT, pmem=extra_mem, 
                                  deps=None, stdout=None, email_addresses=my_email)
        
        command = gatk_VariantEval(ref_genome, merged_tumor_benign_vcf, variant_eval_file_sufix, 
                               interesting_snps_vcf, snpdb_vcf, extra_mem, SINGLE_CORE,
                               path_to_gatk)
        
        jobidrc = jobrunfunc(jobn+'_eval', command, SINGLE_CORE, cwd=None, walltime=WT_SHORT, pmem=extra_mem, 
                                  deps=jobidrc, stdout=None, email_addresses=my_email)
    
    else:
        
        command = gatk_VariantEval2(ref_genome, tumor_vcf, benign_vfc, variant_eval_file_sufix, 
                               interesting_snps_vcf, snpdb_vcf, extra_mem, SINGLE_CORE,
                               path_to_gatk)
        
        jobidrc = jobrunfunc(jobn+'_eval', command, SINGLE_CORE, cwd=None, walltime=WT_SHORT, pmem=extra_mem, 
                                  deps=None, stdout=None, email_addresses=my_email)


    
if __name__ == '__main__':
        
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-r", "--config_file", dest="config_file",
                            help="file with run configuration")
    optionparser.add_option("-a", "--analysis_file", dest="analysis_file",
                            help="file with experiment configuration")
    optionparser.add_option("--snps_hard_filtered", dest="snps_hard_filtered", action="store_true", default=False,
                            help="Determine if it uses the snps obtained by hard filtering or model selection. Default hard_filetered")
    optionparser.add_option("--merge_types", dest="merge_types", action="store_true", default=False,
                            help="Merge vcf files for tumor and benign samples")
    
    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)
    optionparser.add_option("-p", "--processes", type=int, dest="num_processors", default=1)

    (options, args) = optionparser.parse_args()    

    config = ExomePipelineConfig()
    config.from_xml(options.config_file)
    analysis = ExomeAnalysisConfig()
    analysis.from_xml(options.analysis_file, config.output_dir)


    if not (options.local ^ options.cluster):
        print "Must set either --local or --cluster to run job"
    if options.local:
        jobrunfunc = run_local
    elif options.cluster:
        jobrunfunc = qsub
        
    gatk_VariantEvaluation(analysis, config, options.num_processors, jobrunfunc, options.snps_hard_filtered, options.merge_types)
        
    
        
    

